from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import FAISS
import logging
import os
import pickle
import faiss
import tempfile

def vectorize_documents(documents, kon_name, emb_model):
    try:
        # 支持的阿里云Embedding模型
        supported_models = {
            "text-embedding-v1": "embedding-v1",
            "text-embedding-v2": "embedding-v2",
            "text-embedding-async-001": "embedding-async"
        }

        # 根据选择的模型设置参数
        if emb_model in supported_models:
            model_name = supported_models[emb_model]
        else:
            # 默认使用 text-embedding-v2（推荐）
            model_name = "text-embedding-v2"

        print(f"开始创建Embeddings对象，使用模型: {model_name}")

        # 创建Embeddings对象
        embeddings = DashScopeEmbeddings(
            model=model_name,
            dashscope_api_key="sk-c9b8659683a541bfaa8580448ca67766"
        )

        print(f"开始向量化 {len(documents)} 个文档")

        # 检查文档是否为空
        if not documents:
            raise ValueError("没有文档需要向量化")

        # 检查文档内容
        for i, doc in enumerate(documents):
            if not hasattr(doc, 'page_content') or not doc.page_content:
                logging.warning(f"文档 {i} 内容为空")

        vectorstore = FAISS.from_documents(documents, embeddings)

        # 获取 FAISS 索引
        index = vectorstore.index

        # 使用临时文件存储索引
        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
            temp_file_path = temp_file.name

        # 将索引写入临时文件
        faiss.write_index(index, temp_file_path)

        # 读取临时文件内容到内存
        with open(temp_file_path, 'rb') as f:
            faiss_index_data = f.read()

        # 删除临时文件
        os.unlink(temp_file_path)

        # 安全获取文档数量
        n_vectors = len(vectorstore.docstore._dict)  # 直接访问 docs 属性

        index_info = {
            "dimension": index.d,
            "n_vectors": n_vectors,
            "embedding_model": model_name
        }

        # 序列化索引信息为 pickle 数据
        faiss_index_data = faiss.serialize_index(vectorstore.index).tobytes()
        pkl_data = vectorstore.serialize_to_bytes()  # 整库序列化
        return faiss_index_data, pkl_data

    except Exception as e:
        logging.error(f"向量化文档失败: {str(e)}", exc_info=True)
        raise Exception(f"向量化文档失败: {str(e)}")
